{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp tst_net_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "# from nbdev.showdoc import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import torch.nn as nn\n",
    "import sys, torch\n",
    "from functools import partial\n",
    "from collections import OrderedDict\n",
    "\n",
    "from model_constructor.layers import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "act_fn = nn.ReLU(inplace=True)\n",
    "\n",
    "def init_cnn(m):\n",
    "    if getattr(m, 'bias', None) is not None: nn.init.constant_(m.bias, 0)\n",
    "    if isinstance(m, (nn.Conv2d,nn.Linear)): nn.init.kaiming_normal_(m.weight)\n",
    "    for l in m.children(): init_cnn(l)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, expansion, ni, nh, stride=1, sa=False,\n",
    "                 conv_layer=ConvLayer, act_fn=act_fn, zero_bn=True,\n",
    "                 pool=nn.AvgPool2d(2, ceil_mode=True), sym=False):\n",
    "        super().__init__()\n",
    "        nf,ni = nh*expansion,ni*expansion\n",
    "        layers  = [(f\"conv_0\", conv_layer(ni, nh, 3, stride=stride)),\n",
    "                   (f\"conv_1\", conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False))\n",
    "        ] if expansion == 1 else [\n",
    "                   (f\"conv_0\",conv_layer(ni, nh, 1)),\n",
    "                   (f\"conv_1\",conv_layer(nh, nh, 3, stride=stride)),\n",
    "                   (f\"conv_2\",conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False))\n",
    "        ]\n",
    "        if sa: layers.append(('sa', SimpleSelfAttention(nf,ks=1,sym=sym)))\n",
    "        self.convs = nn.Sequential(OrderedDict(layers))\n",
    "        self.pool = noop if stride==1 else pool\n",
    "        self.idconv = noop if ni==nf else conv_layer(ni, nf, 1, act=False)\n",
    "        self.act_fn =act_fn\n",
    "\n",
    "    def forward(self, x): return self.act_fn(self.convs(x) + self.idconv(self.pool(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# class version\n",
    "class _ResBlock(Constructor):\n",
    "    def __init__(self, expansion=1, conv_layer=ConvLayer, act_fn=act_fn, zero_bn=True,\n",
    "                 pool=nn.AvgPool2d(2, ceil_mode=True), sym=False):\n",
    "        super().__init__()\n",
    "        self.__dict__.update(locals())\n",
    "#         self.__dict__.update(self.__dict__.pop('kwargs'))\n",
    "        \n",
    "        \n",
    "    def __call__(self,  ni, nh, stride=1, sa=False):\n",
    "        return ResBlock(self.expansion, ni,nh,stride,sa, \n",
    "                self.conv_layer, self.act_fn, self.zero_bn, self.pool, self.sym)\n",
    "    def __getattr__(self, k): \n",
    "        if hasattr(self, '_model'):\n",
    "            return getattr(self._model, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "block = _ResBlock()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "_ResBlock"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "'bool' object is not callable",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-118-8da458ea545e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mblock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m64\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m64\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-114-d8760d9411bc>\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, ni, nh, stride, sa)\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m  \u001b[0mni\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstride\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msa\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m         return ResBlock(self.expansion, ni,nh,stride,sa, \n\u001b[0;32m---> 13\u001b[0;31m                 self.conv_layer, self.act_fn, self.zero_bn, self.pool, self.sym)\n\u001b[0m\u001b[1;32m     14\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__getattr__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'_model'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-4-65c97fae01da>\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, expansion, ni, nh, stride, conv_layer, act_fn, zero_bn, pool, sa, sym)\u001b[0m\n\u001b[1;32m      8\u001b[0m         layers  = [(f\"conv_0\", conv_layer(ni, nh, 3, stride=stride)),\n\u001b[1;32m      9\u001b[0m                    \u001b[0;34m(\u001b[0m\u001b[0;34mf\"conv_1\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconv_layer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mzero_bn\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mzero_bn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mact\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m         \u001b[0;34m]\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mexpansion\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m                    \u001b[0;34m(\u001b[0m\u001b[0;34mf\"conv_0\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mconv_layer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mni\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m                    \u001b[0;34m(\u001b[0m\u001b[0;34mf\"conv_1\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mconv_layer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnh\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstride\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mstride\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: 'bool' object is not callable"
     ]
    }
   ],
   "source": [
    "block(64,64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (sa): SimpleSelfAttention(\n",
       "      (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)\n",
       "    )\n",
       "  )\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ResBlock(1,64,64,sa=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (act_fn): LeakyReLU(negative_slope=0.01)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ResBlock(2,64,64,act_fn=nn.LeakyReLU())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NewResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# # Still no name - just New block YET!\n",
    "# class NewResBlock(nn.Module):\n",
    "#     def __init__(self, expansion, ni, nh, stride=1, \n",
    "#                  conv_layer=ConvLayer, act_fn=act_fn,\n",
    "#                  pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False, zero_bn=True):\n",
    "#         super().__init__()\n",
    "#         nf,ni = nh*expansion,ni*expansion\n",
    "#         self.reduce = noop if stride==1 else pool\n",
    "#         layers  = [(f\"conv_0\", conv_layer(ni, nh, 3, stride=1)), # stride 1 !!!\n",
    "#                    (f\"conv_1\", conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False))\n",
    "#         ] if expansion == 1 else [\n",
    "#                    (f\"conv_0\",conv_layer(ni, nh, 1)),\n",
    "#                    (f\"conv_1\",conv_layer(nh, nh, 3, stride=1)), # stride 1 !!!\n",
    "#                    (f\"conv_2\",conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False)) ### act!!!\n",
    "#         ]\n",
    "#         if sa: layers.append(('sa', SimpleSelfAttention(nf,ks=1,sym=sym)))\n",
    "#         self.convs = nn.Sequential(OrderedDict(layers))\n",
    "#         self.idconv = noop if ni==nf else conv_layer(ni, nf, 1, act=False)\n",
    "#         self.merge =act_fn\n",
    "\n",
    "#     def forward(self, x): \n",
    "#         o = self.reduce(x)\n",
    "#         return self.merge(self.convs(o) + self.idconv(o))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# Still no name - just New block YET!\n",
    "class NewResBlock(nn.Module):\n",
    "    def __init__(self, expansion, ni, nh, stride=1, \n",
    "                 conv_layer=ConvLayer, act_fn=act_fn,\n",
    "                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False, zero_bn=True):\n",
    "        super().__init__()\n",
    "        nf,ni = nh*expansion,ni*expansion\n",
    "        self.reduce = noop if stride==1 else pool\n",
    "        layers  = [(f\"conv_0\", conv_layer(ni, nh, 3, stride=1)), # stride 1 !!!\n",
    "                   (f\"conv_1\", conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False))\n",
    "        ] if expansion == 1 else [\n",
    "                   (f\"conv_0\",conv_layer(ni, nh, 1)),\n",
    "                   (f\"conv_1\",conv_layer(nh, nh, 3, stride=1)), # stride 1 !!!\n",
    "                   (f\"conv_2\",conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False)) ### act!!!\n",
    "        ]\n",
    "        if sa: layers.append(('sa', SimpleSelfAttention(nf,ks=1,sym=sym)))\n",
    "        self.convs = nn.Sequential(OrderedDict(layers))\n",
    "        self.idconv = noop if ni==nf else conv_layer(ni, nf, 1, act=False)\n",
    "        self.merge =act_fn\n",
    "\n",
    "    def forward(self, x): \n",
    "        o = self.reduce(x)\n",
    "        return self.merge(self.convs(o) + self.idconv(o))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (sa): SimpleSelfAttention(\n",
       "      (conv): Conv1d(64, 64, kernel_size=(1,), stride=(1,), bias=False)\n",
       "    )\n",
       "  )\n",
       "  (merge): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "bl = NewResBlock(1,64,64,sa=True)\n",
    "bl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 64, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "bs_test = 16\n",
    "xb = torch.randn(bs_test, 64, 32, 32)\n",
    "y = bl(xb)\n",
    "print(y.shape)\n",
    "assert y.shape == torch.Size([bs_test, 64, 32, 32]), f\"size\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewResBlock(\n",
       "  (reduce): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (idconv): ConvLayer(\n",
       "    (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): LeakyReLU(negative_slope=0.01)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "bl = NewResBlock(4,64,128,stride=2,act_fn=nn.LeakyReLU())\n",
    "bl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 512, 16, 16])\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "bs_test = 16\n",
    "xb = torch.randn(bs_test, 256, 32, 32)\n",
    "y = bl(xb)\n",
    "print(y.shape)\n",
    "assert y.shape == torch.Size([bs_test, 512, 16, 16]), f\"size\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# class Consructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Constructor():\n",
    "    def __init__(self, **kwargs): \n",
    "        self._name = type(self).__name__\n",
    "        self.__dict__.update(locals())\n",
    "        self.__dict__.update(self.__dict__.pop('kwargs'))\n",
    "        \n",
    "    def __repr__(self): return f\"{self._name}\"  \n",
    "    \n",
    "    def register(self, model):\n",
    "            self._model = model\n",
    "            for attr in self.__dict__.copy():\n",
    "                if not attr.startswith('_') and not attr=='self':\n",
    "                    if not hasattr(model, attr):\n",
    "                        print(f\"setting {attr}\")\n",
    "                        setattr(model, attr, self.__dict__[attr])\n",
    "                    self.__dict__.pop(attr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons = Constructor(x=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Constructor"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_name': 'Constructor', 'self': Constructor, 'x': 1}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.__dict__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## tsts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons_2 = Constructor(bn_1st=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "constr_base"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'constr_base'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons._name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bn_1st': False, '_name': 'constr_base'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'base not found aaa'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.aaa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.Constructor at 0x7f2a20596f10>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class _ConvL(Constructor):\n",
    "#     def __init__(self, parent=cons, conv_layer=ConvLayer, \n",
    "    def __init__(self, parent, conv_layer=ConvLayer, \n",
    "                 act_fn=nn.ReLU(inplace=True), bn_1st=True):\n",
    "        super().__init__()\n",
    "#         in_args = locals()\n",
    "#         print(in_args)\n",
    "#         del in_args['self']\n",
    "#         del in_args['parent']\n",
    "#         del in_args['__class__']\n",
    "#         print(in_args)\n",
    "#         self.parent = parent\n",
    "        \n",
    "#         for arg in in_args:\n",
    "#             if not hasattr(self.parent, arg):\n",
    "#                 setattr(self.parent, arg, in_args[arg])\n",
    "        \n",
    "#     def __getattr__(self, k): return getattr(self.parent, k)\n",
    "# #         act_fn=act_fn, bn_1st=True, \n",
    "# #             padding=None, bias=False, groups=1\n",
    "# #         pass\n",
    "\n",
    "#     def __call__(self, ni, nf, ks=3, stride=1, \n",
    "#             act=True,  bn_layer=True, zero_bn=False):\n",
    "#         return self.conv_layer(ni, nf, ks=3, stride=1, \n",
    "# #         ConvLayer(ni, nf, ks=3, stride=1, \n",
    "# #         NewConvLayer(ni, nf, ks=3, stride=1, \n",
    "#             act=True,  bn_layer=True, zero_bn=False)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l = _ConvL(cons)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<bound method Constructor.__getattr__ of _ConvL>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.__getattr__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'self': _ConvL, '_name': '_ConvL'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l.parent = cons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.Constructor at 0x7f2a1afa3c50>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bn_1st': False, '_name': 'constr_base'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.parent.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.Constructor at 0x7f2a1afa3c50>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.parent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'base not found aaa'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.aaa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.Constructor at 0x7f2a1afa3c50>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.parent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "getattr(cons, 'bn_1st')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "getattr(l.parent, 'bn_1st')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "getattr(l, 'bn_1st')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<bound method _ConvL.__getattr__ of <__main__._ConvL object at 0x7f2a1b702c10>>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.__getattr__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l.__getattr__ = lambda _, k: getattr(cons_2, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "__getattribute__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'base not found aaa'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.aaa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons_2.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'qqq'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "getattr(cons_2, 'aaa')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons_2.aaa = 'qqq'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.<lambda>(_, k)>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l.__getattr__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvLayer(\n",
       "  (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l(64,128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvLayer(\n",
       "  (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ConvLayer(64,128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'bn_1st': False, '_name': 'constr_base'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hasattr(cons, 'act_fn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LeakyReLU(negative_slope=0.01)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.act_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NewConvLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# class NewConvLayer(Constructor):\n",
    "#     \"\"\"Basic conv layers block\"\"\"\n",
    "#     def __init__(self, model, act_fn=act_fn, bn_1st=True, \n",
    "#             padding=None, bias=False, groups=1, norm=nn.BatchNorm2d, **kwargs):\n",
    "# #         self.__dict__ = locals()\n",
    "#         in_args = locals()\n",
    "# #         print(in_args)\n",
    "#         del in_args['self']\n",
    "#         del in_args['model']\n",
    "# #         del in_args['__class__']\n",
    "# #         print(in_args)\n",
    "#         self._model = model\n",
    "#         self._name = 'NewConvLayer'\n",
    "#         for arg in in_args:\n",
    "# #             print(arg)\n",
    "#             if not hasattr(self._model, arg):\n",
    "#                 setattr(self._model, arg, in_args[arg])\n",
    "# #             else: print('exist')\n",
    "        \n",
    "#     def __getattr__(self, k): return getattr(self._model, k)\n",
    "        \n",
    "#     def __call__(self, ni, nf, ks=3, stride=1, \n",
    "#             act=True,  bn_layer=True, zero_bn=False, **kwargs): # todo check kwargs\n",
    "#         padding = ks//2 if self.padding==None else self.padding\n",
    "#         layers = [('conv', nn.Conv2d(ni, nf, ks, stride=stride, padding=padding, \n",
    "#                                      bias=self.bias, groups=self.groups))]\n",
    "#         act_bn = [('act_fn', self.act_fn)] if act else []\n",
    "#         if bn_layer:\n",
    "#             bn = self.norm(nf)\n",
    "#             nn.init.constant_(bn.weight, 0. if zero_bn else 1.) \n",
    "#             act_bn += [('bn', bn)]\n",
    "#         if self.bn_1st: act_bn.reverse()\n",
    "#         layers += act_bn   \n",
    "#         return nn.Sequential(OrderedDict(layers))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#     def register(self, model):\n",
    "#         self._model = model\n",
    "#         for attr in self.__dict__.copy():\n",
    "#             if not attr.startswith('_') and not attr=='self':\n",
    "#                 if not hasattr(model, attr):           \n",
    "#                     setattr(model, attr, self.__dict__[attr])\n",
    "#                 self.__dict__.pop(attr)\n",
    "    \n",
    "#     def __getattr__(self, k): \n",
    "#         if hasattr(self, '_model'): return getattr(self._model, k)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class NewConvLayer(Constructor):\n",
    "    \"\"\"Basic conv layers block\"\"\"\n",
    "    def __init__(self, bn_1st=True, act_fn=act_fn, norm=nn.BatchNorm2d,  \n",
    "            padding=None, bias=False, groups=1,  **kwargs):\n",
    "        super().__init__()\n",
    "        self.__dict__.update(locals())\n",
    "        self.__dict__.update(self.__dict__.pop('kwargs'))\n",
    "    \n",
    "    def __call__(self, ni, nf, ks=3, stride=1, \n",
    "            act=True,  bn_layer=True, zero_bn=False, **kwargs): # todo check kwargs\n",
    "        padding = ks//2 if self.padding==None else self.padding\n",
    "        layers = [('conv', nn.Conv2d(ni, nf, ks, stride=stride, padding=padding, \n",
    "                                     bias=self.bias, groups=self.groups))]\n",
    "        act_bn = [('act_fn', self.act_fn)] if act else []\n",
    "        if bn_layer:\n",
    "            bn = self.norm(nf)\n",
    "            nn.init.constant_(bn.weight, 0. if zero_bn else 1.) \n",
    "            act_bn += [('bn', bn)]\n",
    "        if self.bn_1st: act_bn.reverse()\n",
    "        layers += act_bn   \n",
    "        return nn.Sequential(OrderedDict(layers))\n",
    "    \n",
    "    def __getattr__(self, k): \n",
    "        if hasattr(self, '_model'):\n",
    "            return getattr(self._model, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# _conv_layer = NewConvLayer(cons)\n",
    "_conv_layer = NewConvLayer(z=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewConvLayer"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_name': 'NewConvLayer',\n",
       " 'self': NewConvLayer,\n",
       " 'act_fn': ReLU(inplace=True),\n",
       " 'bn_1st': True,\n",
       " 'padding': None,\n",
       " 'bias': False,\n",
       " 'groups': 1,\n",
       " 'norm': torch.nn.modules.batchnorm.BatchNorm2d,\n",
       " '__class__': __main__.NewConvLayer,\n",
       " 'z': 100}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons = Constructor(bn_1st=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "setting act_fn\n",
      "setting padding\n",
      "setting bias\n",
      "setting groups\n",
      "setting norm\n",
      "setting z\n"
     ]
    }
   ],
   "source": [
    "_conv_layer.register(cons)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer._model.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ReLU(inplace=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer._model.act_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewConvLayer"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_name': 'Constructor',\n",
       " 'self': Constructor,\n",
       " 'bn_1st': False,\n",
       " 'act_fn': ReLU(inplace=True),\n",
       " 'padding': None,\n",
       " 'bias': False,\n",
       " 'groups': 1,\n",
       " 'norm': torch.nn.modules.batchnorm.BatchNorm2d,\n",
       " 'z': 100}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_name': 'NewConvLayer',\n",
       " 'self': NewConvLayer,\n",
       " '__class__': __main__.NewConvLayer,\n",
       " '_model': Constructor}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Constructor"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer._model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ReLU(inplace=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer.act_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_name': 'Constructor',\n",
       " 'self': Constructor,\n",
       " 'bn_1st': False,\n",
       " 'act_fn': ReLU(inplace=True),\n",
       " 'padding': None,\n",
       " 'bias': False,\n",
       " 'groups': 1,\n",
       " 'norm': torch.nn.modules.batchnorm.BatchNorm2d,\n",
       " 'z': 100}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer._model.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "  (act_fn): ReLU(inplace=True)\n",
       "  (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer(64,128,1,2,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons.act_fn = nn.SELU()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "  (act_fn): SELU()\n",
       "  (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer(64,128,1,2,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_name': 'Constructor',\n",
       " 'self': Constructor,\n",
       " 'bn_1st': False,\n",
       " 'act_fn': SELU(),\n",
       " 'padding': None,\n",
       " 'bias': False,\n",
       " 'groups': 1,\n",
       " 'norm': torch.nn.modules.batchnorm.BatchNorm2d,\n",
       " 'z': 100}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cons.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'_name': 'NewConvLayer',\n",
       " 'self': NewConvLayer,\n",
       " '__class__': __main__.NewConvLayer,\n",
       " '_model': Constructor}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons.act_fn = nn.LeakyReLU()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "  (act_fn): LeakyReLU(negative_slope=0.01)\n",
       "  (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer(64,128,1,2,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons.bn_1st = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "  (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (act_fn): LeakyReLU(negative_slope=0.01)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_conv_layer(64,128,1,2,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# v8\n",
    "class Net():\n",
    "    def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000, name='Net'):\n",
    "        super().__init__()\n",
    "        self.name = name\n",
    "        self.c_in, self.c_out,self.expansion,self.layers = c_in,c_out,expansion,layers # todo setter for expansion\n",
    "        self.stem_sizes = [c_in,32,32,64]\n",
    "        self.stem_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "        self.stem_bn_end = False\n",
    "        self.norm = nn.BatchNorm2d\n",
    "        self.act_fn=nn.ReLU(inplace=True)\n",
    "        self.pool = nn.AvgPool2d(2, ceil_mode=True)\n",
    "        self.sa=False\n",
    "        self.bn_1st = True\n",
    "        self.zero_bn=True\n",
    "        self._init_cnn = init_cnn\n",
    "        self.block = ResBlock\n",
    "        self.conv_layer = ConvLayer\n",
    "        \n",
    "    @property\n",
    "    def block_szs(self):\n",
    "        return [64//self.expansion,64,128,256,512] +[256]*(len(self.layers)-4) \n",
    "\n",
    "    @property\n",
    "    def stem(self):\n",
    "        return self._make_stem()\n",
    "    @property\n",
    "    def head(self):\n",
    "        return self._make_head()\n",
    "    @property\n",
    "    def body(self):\n",
    "        return self._make_body()\n",
    "    \n",
    "    def _make_stem(self):\n",
    "        stem = [(f\"conv_{i}\", self.conv_layer(self.stem_sizes[i], self.stem_sizes[i+1], \n",
    "                    stride=2 if i==0 else 1, \n",
    "                    norm=(not self.stem_bn_end) if i==(len(self.stem_sizes)-2) else True,\n",
    "                    act_fn=self.act_fn, bn_1st=self.bn_1st))\n",
    "                for i in range(len(self.stem_sizes)-1)]\n",
    "        stem.append(('stem_pool', self.stem_pool))\n",
    "        if self.stem_bn_end: stem.append(('norm', self.norm(self.stem_sizes[-1])))\n",
    "        return nn.Sequential(OrderedDict(stem))\n",
    "    \n",
    "    def _make_head(self):\n",
    "        head = [('pool', nn.AdaptiveAvgPool2d(1)),\n",
    "                ('flat', Flatten()),\n",
    "                ('fc',   nn.Linear(self.block_szs[-1]*self.expansion, self.c_out))]\n",
    "        return nn.Sequential(OrderedDict(head))\n",
    "    \n",
    "    def _make_body(self):\n",
    "        blocks = [(f\"l_{i}\", self._make_layer(self.block_szs[i], self.block_szs[i+1], l, \n",
    "                        1 if i==0 else 2, self.sa if i==0 else False))\n",
    "                  for i,l in enumerate(self.layers)]\n",
    "        return nn.Sequential(OrderedDict(blocks))\n",
    "    \n",
    "    def _make_layer(self,ni,nf,blocks,stride,sa):\n",
    "        return nn.Sequential(OrderedDict(\n",
    "            [(f\"bl_{i}\", self.block(self.expansion, ni if i==0 else nf, nf, \n",
    "                    stride if i==0 else 1, sa=sa if i==blocks-1 else False,\n",
    "                    conv_layer=self.conv_layer, act_fn=self.act_fn, pool=self.pool,zero_bn=self.zero_bn))\n",
    "              for i in range(blocks)]))\n",
    "    \n",
    "    def __call__(self):\n",
    "        model = nn.Sequential(OrderedDict([\n",
    "            ('stem', self.stem),\n",
    "            ('body', self.body),\n",
    "            ('head', self.head)\n",
    "        ]))\n",
    "        self._init_cnn(model)\n",
    "        model.extra_repr = lambda : f\"model {self.name}\"\n",
    "        return model\n",
    "    def __repr__(self):\n",
    "        return f\" constr {self.name}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Net class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# v9\n",
    "class Net():\n",
    "    def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000, name='Net'):\n",
    "        super().__init__()\n",
    "        self.name = name\n",
    "        self.c_in, self.c_out,self.expansion,self.layers = c_in,c_out,expansion,layers # todo setter for expansion\n",
    "        self.stem_sizes = [c_in,32,32,64]\n",
    "        self.stem_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "        self.stem_bn_end = False\n",
    "        self.norm = nn.BatchNorm2d\n",
    "        self.act_fn=nn.ReLU(inplace=True)\n",
    "        self.pool = nn.AvgPool2d(2, ceil_mode=True)\n",
    "        self.sa=False\n",
    "        self.bn_1st = True\n",
    "        self.zero_bn=True\n",
    "        self._init_cnn = init_cnn\n",
    "        self.block = ResBlock\n",
    "        \n",
    "#         self.conv_layer = ConvLayer\n",
    "        self.conv_layer = NewConvLayer\n",
    "    \n",
    "    @property\n",
    "    def conv_layer(self): return self._conv_layer\n",
    "    @conv_layer.setter\n",
    "    def conv_layer(self, f):\n",
    "        self._conv_layer = f()\n",
    "        self._conv_layer.register(self)\n",
    "        \n",
    "    \n",
    "    @property\n",
    "    def block_szs(self):\n",
    "        return [64//self.expansion,64,128,256,512] +[256]*(len(self.layers)-4) \n",
    "\n",
    "    @property\n",
    "    def stem(self):\n",
    "        return self._make_stem()\n",
    "    @property\n",
    "    def head(self):\n",
    "        return self._make_head()\n",
    "    @property\n",
    "    def body(self):\n",
    "        return self._make_body()\n",
    " \n",
    "    def _block(self, ni, nh, stride=1, sa=False,sym=False):\n",
    "        return self.block(self.expansion, ni, nh, stride, \n",
    "                 self.conv_layer, self.act_fn,\n",
    "                 pool=self.pool, sa=sa,sym=sym, zero_bn=self.zero_bn)\n",
    "    \n",
    "    \n",
    "    def _make_stem(self):\n",
    "        stem = [(f\"conv_{i}\", self._conv_layer(self.stem_sizes[i], self.stem_sizes[i+1], \n",
    "                    stride=2 if i==0 else 1, \n",
    "                    bn_layer=(not self.stem_bn_end) if i==(len(self.stem_sizes)-2) else True,))\n",
    "                for i in range(len(self.stem_sizes)-1)]\n",
    "        stem.append(('stem_pool', self.stem_pool))\n",
    "        if self.stem_bn_end: stem.append(('norm', self.norm(self.stem_sizes[-1])))\n",
    "        return nn.Sequential(OrderedDict(stem))\n",
    "    \n",
    "    def _make_head(self):\n",
    "        head = [('pool', nn.AdaptiveAvgPool2d(1)),\n",
    "                ('flat', Flatten()),\n",
    "                ('fc',   nn.Linear(self.block_szs[-1]*self.expansion, self.c_out))]\n",
    "        return nn.Sequential(OrderedDict(head))\n",
    "    \n",
    "    def _make_body(self):\n",
    "        blocks = [(f\"l_{i}\", self._make_layer(self.block_szs[i], self.block_szs[i+1], l, \n",
    "                        1 if i==0 else 2, self.sa if i==0 else False))\n",
    "                  for i,l in enumerate(self.layers)]\n",
    "        return nn.Sequential(OrderedDict(blocks))\n",
    "    \n",
    "    def _make_layer(self,ni,nf,blocks,stride,sa):\n",
    "        return nn.Sequential(OrderedDict(\n",
    "            [(f\"bl_{i}\", self._block(ni if i==0 else nf, nf, \n",
    "                    stride if i==0 else 1, sa=sa if i==blocks-1 else False))\n",
    "              for i in range(blocks)]))\n",
    "    \n",
    "    def __call__(self):\n",
    "        model = nn.Sequential(OrderedDict([\n",
    "            ('stem', self.stem),\n",
    "            ('body', self.body),\n",
    "            ('head', self.head)\n",
    "        ]))\n",
    "        self._init_cnn(model)\n",
    "        model.extra_repr = lambda : f\"model {self.name}\"\n",
    "        return model\n",
    "    def __repr__(self):\n",
    "        return f\" constr {self.name}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "setting padding\n",
      "setting bias\n",
      "setting groups\n"
     ]
    }
   ],
   "source": [
    "model  = Net()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ReLU(inplace=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.conv_layer.act_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewConvLayer"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.conv_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.conv_layer = NewConvLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.conv_layer(64,128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       " constr Net"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[64, 64, 128, 256, 512]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.block_szs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv_0): Sequential(\n",
       "    (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): ReLU(inplace=True)\n",
       "  )\n",
       "  (conv_1): Sequential(\n",
       "    (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): ReLU(inplace=True)\n",
       "  )\n",
       "  (conv_2): Sequential(\n",
       "    (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): ReLU(inplace=True)\n",
       "  )\n",
       "  (stem_pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.stem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 64, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "bs_test = 16\n",
    "xb = torch.randn(bs_test, 3, 128, 128)\n",
    "y = model.stem(xb)\n",
    "print(y.shape)\n",
    "assert y.shape == torch.Size([bs_test, 64, 32, 32]), f\"size\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.bn_1st = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.act_fn =nn.LeakyReLU(inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LeakyReLU(negative_slope=0.01, inplace=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.conv_layer.act_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.sa = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.conv_layer =NewConvLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.conv_layer.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LeakyReLU(negative_slope=0.01, inplace=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.conv_layer.act_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv_0): Sequential(\n",
       "    (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (conv_1): Sequential(\n",
       "    (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (conv_2): Sequential(\n",
       "    (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (stem_pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.stem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.block = NewResBlock\n",
    "model.block = ResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.bn_1st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "  (idconv): ConvLayer(\n",
       "    (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.block(1,64,128,2,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (bl_0): ResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): Sequential(\n",
       "        (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (conv_1): Sequential(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "    (idconv): Sequential(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       "  (bl_1): ResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): Sequential(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (conv_1): Sequential(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.body.l_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): Sequential(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (conv_1): Sequential(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model._block(64,64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (convs): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): ReLU(inplace=True)\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (pool): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "  (act_fn): ReLU(inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.block(1,64,64,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 64, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "bs_test = 16\n",
    "xb = torch.randn(bs_test, 64, 32, 32)\n",
    "y = model.body.l_0(xb)\n",
    "print(y.shape)\n",
    "assert y.shape == torch.Size([bs_test, 64, 32, 32]), f\"size\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.block = NewResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 128, 16, 16])\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "bs_test = 16\n",
    "xb = torch.randn(bs_test, 64, 32, 32)\n",
    "y = model.body.l_1.bl_0(xb)\n",
    "print(y.shape)\n",
    "assert y.shape == torch.Size([bs_test, 128, 16, 16]), f\"size\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewResBlock(\n",
       "  (reduce): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
       "  (convs): Sequential(\n",
       "    (conv_0): Sequential(\n",
       "      (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (conv_1): Sequential(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (idconv): Sequential(\n",
       "    (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model.body.l_1.bl_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model = Net(expansion=4)\n",
    "model.expansion = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.stem_bn_end = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.bn_1st = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv_0): Sequential(\n",
       "    (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       "  (conv_1): Sequential(\n",
       "    (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       "  (conv_2): Sequential(\n",
       "    (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       "  (stem_pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.stem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AvgPool2d(kernel_size=2, stride=2, padding=0)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.pool = nn.MaxPool2d(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewResBlock(\n",
       "  (reduce): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
       "  (convs): Sequential(\n",
       "    (conv_0): Sequential(\n",
       "      (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_1): Sequential(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_2): Sequential(\n",
       "      (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (idconv): Sequential(\n",
       "    (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.body.l_1.bl_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.block = NewResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NewResBlock(\n",
       "  (reduce): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
       "  (convs): Sequential(\n",
       "    (conv_0): Sequential(\n",
       "      (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_1): Sequential(\n",
       "      (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_2): Sequential(\n",
       "      (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "  )\n",
       "  (idconv): Sequential(\n",
       "    (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model.body.l_1.bl_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 512, 10, 10])\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "bs_test = 16\n",
    "xb = torch.randn(bs_test, 256, 32, 32)\n",
    "y = model.body.l_1.bl_0(xb)\n",
    "print(y.shape)\n",
    "# assert y.shape == torch.Size([bs_test, 512, 16, 16]), f\"size\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  model Net\n",
       "  (stem): Sequential(\n",
       "    (conv_0): Sequential(\n",
       "      (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_1): Sequential(\n",
       "      (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (conv_2): Sequential(\n",
       "      (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "    )\n",
       "    (stem_pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "    (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  )\n",
       "  (body): Sequential(\n",
       "    (l_0): Sequential(\n",
       "      (bl_0): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): Sequential(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_1): Sequential(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_2): Sequential(\n",
       "            (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (idconv): Sequential(\n",
       "          (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      )\n",
       "      (bl_1): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): Sequential(\n",
       "            (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_1): Sequential(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_2): Sequential(\n",
       "            (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "          (sa): SimpleSelfAttention(\n",
       "            (conv): Conv1d(256, 256, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "        )\n",
       "        (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (l_1): Sequential(\n",
       "      (bl_0): NewResBlock(\n",
       "        (reduce): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
       "        (convs): Sequential(\n",
       "          (conv_0): Sequential(\n",
       "            (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_1): Sequential(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_2): Sequential(\n",
       "            (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (idconv): Sequential(\n",
       "          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      )\n",
       "      (bl_1): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): Sequential(\n",
       "            (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_1): Sequential(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_2): Sequential(\n",
       "            (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (l_2): Sequential(\n",
       "      (bl_0): NewResBlock(\n",
       "        (reduce): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
       "        (convs): Sequential(\n",
       "          (conv_0): Sequential(\n",
       "            (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_1): Sequential(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_2): Sequential(\n",
       "            (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (idconv): Sequential(\n",
       "          (conv): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      )\n",
       "      (bl_1): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): Sequential(\n",
       "            (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_1): Sequential(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_2): Sequential(\n",
       "            (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (l_3): Sequential(\n",
       "      (bl_0): NewResBlock(\n",
       "        (reduce): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
       "        (convs): Sequential(\n",
       "          (conv_0): Sequential(\n",
       "            (conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_1): Sequential(\n",
       "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_2): Sequential(\n",
       "            (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (idconv): Sequential(\n",
       "          (conv): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      )\n",
       "      (bl_1): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): Sequential(\n",
       "            (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_1): Sequential(\n",
       "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "          )\n",
       "          (conv_2): Sequential(\n",
       "            (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (head): Sequential(\n",
       "    (pool): AdaptiveAvgPool2d(output_size=1)\n",
       "    (flat): Flatten()\n",
       "    (fc): Linear(in_features=2048, out_features=1000, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 1000])\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "bs_test = 16\n",
    "xb = torch.randn(bs_test, 3, 128, 128)\n",
    "y = m(xb)\n",
    "print(y.shape)\n",
    "assert y.shape == torch.Size([bs_test, 1000]), f\"size expected {bs_test}, 1000\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.stem_sizes = [3, 32, 32, 64, 64]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv_0): Sequential(\n",
       "    (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       "  (conv_1): Sequential(\n",
       "    (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       "  (conv_2): Sequential(\n",
       "    (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       "  (conv_3): Sequential(\n",
       "    (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "  )\n",
       "  (stem_pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (norm): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model.stem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (pool): AdaptiveAvgPool2d(output_size=1)\n",
       "  (flat): Flatten()\n",
       "  (fc): Linear(in_features=2048, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model.head"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## xresnet constructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "me = sys.modules[__name__]\n",
    "for n,e,l in [[ 18 , 1, [2,2,2 ,2] ],\n",
    "    [ 34 , 1, [3,4,6 ,3] ],\n",
    "    [ 50 , 4, [3,4,6 ,3] ],\n",
    "    [ 101, 4, [3,4,23,3] ],\n",
    "    [ 152, 4, [3,8,36,3] ],]:\n",
    "    name = f'net{n}'\n",
    "    setattr(me, name, partial(Net, expansion=e, layers=l, name=name))\n",
    "xresnet50      = partial(Net, expansion=4, layers=[3, 4,  6, 3], name='xresnet50')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = xresnet50(c_out=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "( constr xresnet50, 10)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m, m.c_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# end\n",
    "model_constructor\n",
    "by ayasyrev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hide\n",
    "from nbdev.export import *\n",
    "notebook2script()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
